! test_with_hip_omp.f90
module test
   use, intrinsic :: iso_c_binding, only: c_int, c_float, c_ptr, c_size_t, c_loc, c_sizeof, c_intptr_t
   use hip_wrapper, only: hipMalloc, hipFree, hipMemcpyAsync, &
                         hipStreamCreate, hipStreamSynchronize, hipSuccess, &
                         hipMemcpyHostToDevice, hipMemcpyDeviceToHost
   implicit none

    ! 声明内核启动包装函数（需与C/C++实现匹配）
   interface
      subroutine launch_matrix_transpose_kernel(d_output, d_input, width, &
                                             grid_size_x, grid_size_y, &
                                             threads_per_block_x, threads_per_block_y, &
                                             shared_mem, stream, hip_err) &
                                             bind(C, name="launch_matrix_transpose_kernel")
         use, intrinsic :: iso_c_binding, only: c_ptr, c_size_t, c_int, c_intptr_t
         use hip_wrapper, only: hipSuccess
         implicit none
         type(c_ptr), value :: d_output, d_input
         integer(c_int), value :: width, grid_size_x, grid_size_y
         integer(c_int), value :: threads_per_block_x, threads_per_block_y
         integer(c_size_t), value :: shared_mem
         integer(c_size_t), value :: stream 
         integer(c_int), intent(out) :: hip_err
         
      end subroutine
   end interface
    
   contains
    
   subroutine vec_mult(p, v1, v2, N)
      integer, intent(in) :: N  ! Add intent(in) for N
      real    ::  p(N), v1(N), v2(N)
      integer ::  i,j
      integer :: hip_err
      logical :: hip_available
      integer :: matrix_width
      parameter (matrix_width = 1024)
      real(c_float), allocatable :: matrix(:), transpose_matrix(:), cpu_transpose_matrix(:)
      type(c_ptr) :: d_input, d_output  ! 移除不必要的指针声明
      integer(c_int) :: threads_per_block_x, threads_per_block_y
      integer(c_int) :: grid_size_x, grid_size_y
      integer(c_size_t) :: stream 
      integer(c_size_t) :: shared_mem
      integer(c_int) :: copy_kind 
      type(c_ptr) :: dst_ptr
      
      ! 初始化向量
      call init(v1, v2, N)

      ! OpenMP 向量乘法部分
      p = 0
      !$omp target data device(2) map(to:v1,v2) map(from:p)
         !$omp target device(1)
            do i=1,N
               p(i) = v1(i) * v2(i)
            end do
         !$omp end target
      !$omp end target data
      if(p(1)/=0)then
         write(*,*)"p0 = ",p(1)
      endif

      p = 0
      !$omp target data device(2) map(to:v1,v2) map(from:p)
         !$omp target device(2)
            do i=1,N
               p(i) = v1(i) * v2(i)
            end do
         !$omp end target
      !$omp end target data
      if(p(1)/=6.0)then
         write(*,*)"p1 = ",p(1)
      endif


      p = 0
      !$omp target data device(2) map(to:v1,v2) map(from:p)
         !$omp target device(2)
            do i=1,N
               p(i) = v1(i) * v2(i)
            end do
         !$omp end target
         !$omp target update from(p) device(2)
      !$omp end target data
      if(p(1)/=6.0)then
         write(*,*)"p2 = ",p(1)
      endif

      p = 0
      !$omp target data device(2) map(to:v1,v2) map(from:p)
         !$omp target device(2)
            do i=1,N
               p(i) = v1(i) * v2(i)
            end do
         !$omp end target
         !$omp target update from(p) device(2)
         if(p(1)/=6.0)then
            write(*,*)"p3 = ",p(1)
         endif
      !$omp end target data
      
      ! HIP 矩阵转置部分
      write(*,*) "开始执行 HIP 矩阵转置..."
      
      ! 实际检查HIP环境可用性
      hip_available = check_hip_availability(hip_err)
      if (.not. hip_available) then
          write(*,*) "HIP 环境不可用，跳过 HIP 部分"
          return
      end if
      
      ! 分配主机内存
      allocate(matrix(matrix_width*matrix_width))
      allocate(transpose_matrix(matrix_width*matrix_width))
      allocate(cpu_transpose_matrix(matrix_width*matrix_width))
      
      ! 初始化矩阵数据
      call init_matrix(matrix, matrix_width)
      
      ! 执行 HIP 矩阵转置
      call hip_matrix_transpose(matrix, transpose_matrix, matrix_width, hip_err)
      
      if (hip_err == hipSuccess) then
         ! 执行 CPU 参考转置
         call matrix_transpose_cpu(cpu_transpose_matrix, matrix, matrix_width)
         
         ! 验证结果
         call verify_results(transpose_matrix, cpu_transpose_matrix, matrix_width*matrix_width, hip_err)
         
         if (hip_err == 0) then
            write(*,*) "HIP 矩阵转置验证通过!"
         else
            write(*,*) "HIP 矩阵转置验证失败，错误数: ", hip_err
         end if
      else
         write(*,*) "HIP 矩阵转置执行错误: ", hip_err
      end if
      
      ! 释放内存
      deallocate(matrix)
      deallocate(transpose_matrix)
      deallocate(cpu_transpose_matrix)
          
   end subroutine
    
   subroutine init(v1, v2, N)
      integer,intent(in)::N
      real,intent(inout)::v1(N), v2(N)
      integer :: i
      v1 = 2.0
      v2 = 3.0
   end subroutine 
   
   subroutine init_matrix(matrix, width)
      real(c_float), intent(out) :: matrix(:)
      integer(c_int), intent(in) :: width
      integer :: i
      do i = 1, size(matrix)
         matrix(i) = real(i-1, c_float) * 10.0_c_float
      end do
   end subroutine
    
   subroutine matrix_transpose_cpu(output, input, width)
      real(c_float), intent(out) :: output(:)
      real(c_float), intent(in) :: input(:)
      integer(c_int), intent(in) :: width
      integer :: i, j
      do j = 1, width
         do i = 1, width
            output(i + (j-1)*width) = input(j + (i-1)*width)
         end do
      end do
   end subroutine
    
   subroutine verify_results(result, reference, size, errors)
      real(c_float), intent(in) :: result(:), reference(:)
      integer(c_int), intent(in) :: size
      integer(c_int), intent(out) :: errors
      integer :: i
      real(c_float) :: eps
      eps = 1.0E-6_c_float
      errors = 0
      do i = 1, size
         if (abs(result(i) - reference(i)) > eps) then
            errors = errors + 1
         end if
      end do
   end subroutine
    
   logical function check_hip_availability(hip_err)
      use hip_wrapper, only: hipInit, hipSuccess
      integer(c_int), intent(out) :: hip_err
      integer(c_int) :: version
      hip_err = hipSuccess
      check_hip_availability = .false.
      
      ! 实际初始化HIP运行时
      version = 0
      hip_err = hipInit(0)
      if (hip_err == hipSuccess) then
         check_hip_availability = .true.
      else
         write(*,*) "HIP初始化错误: ", hip_err
      end if
   end function
   
     
   subroutine hip_matrix_transpose(input, output, width, hip_err)
      use, intrinsic :: iso_c_binding, only: c_ptr, c_size_t, c_int, c_loc, c_sizeof, c_intptr_t
      use hip_wrapper, only: hipSuccess, hipMemcpyHostToDevice, hipMemcpyDeviceToHost, &
                           hipMalloc, hipFree, hipMemcpyAsync, hipStreamCreate, &
                           hipStreamSynchronize, hipStreamDestroy
      implicit none
      
      real(c_float), intent(in) :: input(:)
      real(c_float), intent(out) :: output(:)
      integer(c_int), intent(in) :: width
      integer(c_int), intent(out) :: hip_err
      type(c_ptr) :: d_input, d_output
      integer(c_int) :: threads_per_block_x, threads_per_block_y
      integer(c_int) :: grid_size_x, grid_size_y
      integer(c_size_t) :: stream 
      integer(c_size_t) :: shared_mem
      type(c_ptr) :: dst_ptr
      
      ! 初始化错误码
      hip_err = hipSuccess
      
      ! 设置线程块和网格维度
      threads_per_block_x = 16
      threads_per_block_y = 16
      grid_size_x = (width + threads_per_block_x - 1) / threads_per_block_x
      grid_size_y = (width + threads_per_block_y - 1) / threads_per_block_y
      shared_mem = threads_per_block_x * threads_per_block_y * c_sizeof(c_float)
      
      ! 创建流 (直接赋值给 stream)
      hip_err = hipStreamCreate(stream)
      if (hip_err /= hipSuccess) return
      
      ! 分配设备内存
      hip_err = hipMalloc(d_input, size(input)*c_sizeof(c_float))
      if (hip_err /= hipSuccess) goto 99
      
      hip_err = hipMalloc(d_output, size(output)*c_sizeof(c_float))
      if (hip_err /= hipSuccess) goto 99
      
      ! 主机到设备内存传输
      hip_err = hipMemcpyAsync(d_input, c_loc(input(1)), &
                              int(size(input)*c_sizeof(c_float), c_size_t), &
                              hipMemcpyHostToDevice, stream)
      if (hip_err /= hipSuccess) goto 99
      
      ! 启动内核
      call launch_matrix_transpose_kernel(d_output, d_input, width, &
                                       grid_size_x, grid_size_y, &
                                       threads_per_block_x, threads_per_block_y, &
                                       shared_mem, stream, hip_err)
      if (hip_err /= hipSuccess) goto 99
      
      ! 设备到主机内存传输
      dst_ptr = c_loc(output(1))
      hip_err = hipMemcpyAsync(dst_ptr, d_output, &
                              int(size(output)*c_sizeof(c_float), c_size_t), &
                              hipMemcpyDeviceToHost, stream)
      if (hip_err /= hipSuccess) goto 99
      
      ! 同步流
      hip_err = hipStreamSynchronize(stream)
      if (hip_err /= hipSuccess) goto 99
      
   99  continue
      ! 释放设备内存
      hip_err = hipFree(d_input)
      hip_err = hipFree(d_output)
      
      ! 销毁流
      hip_err = hipStreamDestroy(stream)
   end subroutine

end module


program main
use test
   integer, parameter :: N=1024
   real    ::  p(N), v1(N), v2(N)
   
   write(*,*) "开始执行程序..."
   call vec_mult(p, v1, v2, N)
   write(*,*) "程序执行完毕"

end program